#!/usr/bin/env python3
"""
LoRA训练器节点基类 - 提取共同的代码逻辑
"""

import os
import sys
import json
import threading
from pathlib import Path
from abc import ABC, abstractmethod

# 导入ComfyUI模块
try:
    import execution
    from server import PromptServer
    from comfy.model_management import interrupt_current_processing
except ImportError:
    interrupt_current_processing = None

# 导入训练器工具
from .lora_trainer_utils.constants import DEFAULT_VAE_NAME

import folder_paths

# 确保exit函数可用
if not hasattr(__builtins__, 'exit'):
    __builtins__['exit'] = sys.exit

# 全局中断标志
_training_interrupted = False
_training_thread = None

def set_training_interrupted():
    """设置训练中断标志"""
    global _training_interrupted
    _training_interrupted = True
    print("[INFO] 训练中断信号已设置")
    
    if interrupt_current_processing:
        try:
            interrupt_current_processing()
        except Exception as e:
            print(f"[WARNING] ComfyUI中断失败: {e}")
    return True

def check_interrupt():
    """检查是否有中断信号"""
    global _training_interrupted
    
    if _training_interrupted:
        return True
        
    try:
        import comfy.model_management as mm
        if hasattr(mm, 'processing_interrupted') and mm.processing_interrupted():
            return True
    except Exception:
        pass
        
    return False

class BaseLoRATrainerNode(ABC):
    """LoRA训练器节点基类"""
    
    def __init__(self, trainer_module_name, trainer_class_name):
        """初始化训练器节点
        Args:
            trainer_module_name: 训练器模块名称
            trainer_class_name: 训练器类名称
        """
        self.trainer_module_name = trainer_module_name
        self.trainer_class_name = trainer_class_name
        self.trainer = None  # 延迟创建训练器实例
        self.output_dir = "./output"
        self.training_thread = None
        self.stop_training = False
        self.current_node_id = None
        self.progress = {"current": 0, "total": 0, "status": "未开始"}
        
        # 添加ComfyUI中断监听
        self.interrupt_check_thread = None
        self.is_monitoring_interrupt = False
        
        # 导入ComfyUI中断检查函数
        try:
            import comfy.model_management as mm
            self.processing_interrupted = mm.processing_interrupted
            self.COMFYUI_INTERRUPT_AVAILABLE = True
            print("[DEBUG] ComfyUI中断检查函数导入成功")
        except ImportError as e:
            self.processing_interrupted = None
            self.COMFYUI_INTERRUPT_AVAILABLE = False
            print(f"[DEBUG] ComfyUI中断检查函数导入失败: {e}")
        
        print(f"[DEBUG] BaseLoRATrainerNode初始化完成，模块: {trainer_module_name}, 类: {trainer_class_name}, 中断功能可用: {self.COMFYUI_INTERRUPT_AVAILABLE}")
    
    def _create_trainer_instance(self):
        """创建训练器实例"""
        try:
            # 直接创建训练器实例，无需中间层
            from .lora_trainer_utils.base_trainer_api import SimpleTrainerAPI
            return SimpleTrainerAPI(self.trainer_module_name, self.trainer_class_name)
        except Exception as e:
            print(f"[ERROR] 创建训练器实例失败: {e}")
            raise
    
    def start_interrupt_monitoring(self):
        """开始监听中断信号"""
        if not self.processing_interrupted or self.is_monitoring_interrupt:
            return
            
        self.is_monitoring_interrupt = True
        self.interrupt_check_thread = threading.Thread(
            target=self._interrupt_monitor_loop,
            daemon=True
        )
        self.interrupt_check_thread.start()
    
    def stop_interrupt_monitoring(self):
        """停止监听中断信号"""
        if not self.is_monitoring_interrupt:
            return
        
        self.is_monitoring_interrupt = False
        self.interrupt_check_thread = None
    
    def _interrupt_monitor_loop(self):
        """中断监听循环"""
        import time
        
        while self.is_monitoring_interrupt:
            try:
                if self.processing_interrupted():
                    if self.trainer and hasattr(self.trainer, 'stop_training'):
                        self.trainer.stop_training()
                    break
                time.sleep(0.1)  # 每100ms检查一次
            except Exception as e:
                print(f"[ERROR] 中断监听出错: {e}")
                break
    
    def update_progress(self, current_step, total_steps, loss=None, status="训练中"):
        """更新训练进度 - 模板方法"""
        print(f"[DEBUG] 🔄 BaseLoRATrainerNode.update_progress被调用: step={current_step}, total={total_steps}, loss={loss}, status={status}, node_id={self.current_node_id}")
        
        progress_data = {
            "node_id": self.current_node_id,
            "current_step": current_step,
            "total_steps": total_steps,
            "loss": loss,
            "status": status
        }
        
        print(f"[DEBUG] 📤 准备发送进度事件: {progress_data}")
        
        # 调用子类实现的发送方法
        self._send_progress_event(progress_data)
        
        print(f"[DEBUG] ✅ 进度事件发送完成")
    
    @abstractmethod
    def _send_progress_event(self, progress_data: dict):
        """发送进度事件 - 子类必须实现此方法
        
        Args:
            progress_data (dict): 进度数据字典，包含以下字段：
                - node_id: 节点ID
                - current_step: 当前步数
                - total_steps: 总步数
                - loss: 损失值（可选）
                - status: 训练状态
        
        Note:
            子类必须实现此方法，用于发送特定类型的进度事件。
            可以自定义发送逻辑，如不同的事件名、额外的处理等。
        """
        pass
    
    @classmethod
    def IS_CHANGED(cls, **kwargs):
        """强制节点重新执行"""
        return float("nan")
    
    def _validate_training_data_path(self, train_data_dir):
        """验证训练数据路径"""
        if not train_data_dir:
            return "错误：训练数据目录路径不能为空"
        
        if not os.path.exists(train_data_dir):
            return f"错误：训练数据目录不存在: {train_data_dir}"
        
        if not os.path.isdir(train_data_dir):
            return f"错误：训练数据路径不是目录: {train_data_dir}"
        
        # 检查是否有图片文件
        image_extensions = {'.jpg', '.jpeg', '.png', '.webp', '.bmp'}
        has_images = False
        
        for root, dirs, files in os.walk(train_data_dir):
            for file in files:
                if any(file.lower().endswith(ext) for ext in image_extensions):
                    has_images = True
                    break
            if has_images:
                break
        
        if not has_images:
            return f"错误：训练数据目录中没有找到图片文件: {train_data_dir}"
        
        return "验证通过"
    
    def train_lora(self, training_params, node_id=None):
        """训练LoRA模型 - 通用训练方法"""
        try:
            # 保存node_id
            self.current_node_id = node_id
            
            # 每次训练时创建新的训练器实例
            print(f"[INFO] 创建新的训练器实例: {self.trainer_module_name}.{self.trainer_class_name}")
            self.trainer = self._create_trainer_instance()
            
                    # 从training_params中提取基础参数
            base_model = training_params.get("pretrained_model_name_or_path", "")
            train_data_dir = training_params.get("train_data_dir", "")
            output_dir = training_params.get("output_dir", "")
            output_name = training_params.get("output_name", "trained_flux_lora")
            config_file = training_params.get("config_file", "")
            # 验证路径
            status_msg = self._validate_training_data_path(train_data_dir)
            if "错误" in status_msg:
                return {
                    "ui": {"error": status_msg},
                    "result": (status_msg, "", "")
                }
            
           
            if not os.path.exists(base_model):
                return {
                    "ui": {"error": f"错误：找不到基础模型: {base_model}"},
                    "result": (f"错误：找不到基础模型: {base_model}", "", "")
                }
            
            # 处理输出目录
            output_dir = os.path.abspath(output_dir)
            os.makedirs(output_dir, exist_ok=True)
        
            
        
            
            print(f"[INFO] 🎯 开始训练LoRA: {output_name}")
            print(f"[INFO] 📁 基础模型: {base_model}")
            print(f"[INFO] 📁 训练数据: {train_data_dir}")
            print(f"[INFO] 📁 输出目录: {output_dir}")
            
            # 开始训练
            self.start_interrupt_monitoring()
                
            try:
                result = self.trainer.train(training_params, progress_callback=self.update_progress)
                
                # 获取输出文件路径
                output_path = os.path.join(output_dir, f"{output_name}.safetensors")
                
                return {
                    "ui": {
                        "status": "训练完成",
                        "output_path": output_path
                    },
                    "result": (
                        "训练完成",  # 训练状态
                        output_path, # 输出路径
                        "训练日志"   # 训练日志
                    )
                }
            finally:
                self.stop_interrupt_monitoring()
                # 清理训练器实例
                self.trainer = None
            
        except Exception as e:
            print(f"[ERROR] 训练失败: {e}")
            import traceback
            traceback.print_exc()
            # 清理训练器实例
            self.trainer = None
            return {
                "ui": {"error": str(e)},
                "result": (f"训练失败: {str(e)}", "", "")
            } 
    
